from .base_strategy import BaseStrategy
import backtrader as bt
import backtrader.indicators as btind
import indicator.MyTTIndicator as Mytt
import indicator.myzhibiao as mmm


###
# 10日均线回归策略（但是写错了，实际上是均线追随策略，是趋势策略）
# 当价格偏离均线过多，且，出现了上上影线时，则意味着当前价格处于超卖。可以进场卖出
# 
# #
class MyStrategy(BaseStrategy):
    params = (
        ('ema_period', 10),  # EMA10
        ('stop_loss_pct', 2),  # 20% 止损
        ('take_profit_pct', 0),  # 止盈百分比，这里设置为0，因为您提到的是回到EMA均线时止盈
        ('risk_factor',0.5) # 开仓资金占比
    )

 
    def __init__(self):
        super().__init__()

        self.ema = btind.EMA(self.data.close, period=self.p.ema_period)
    
        # 计算K线实体和影线
        self.body = abs(self.data.close - self.data.open)
        self.upper_shadow = bt.indicators.If(
            self.data.close > self.data.open,
            self.data.high - self.data.close,
            self.data.high - self.data.open
        )
        self.lower_shadow = bt.indicators.If(
            self.data.close > self.data.open,
            self.data.open - self.data.low,
            self.data.close - self.data.low
        )
        
        # 判断K线实体是否和均线相交
        self.above_ema = bt.indicators.And(
            self.data.close > self.ema,
            self.data.open > self.ema
        )
        self.below_ema = bt.indicators.And(
            self.data.close < self.ema,
            self.data.open < self.ema
        )
        self.cross_ema = bt.indicators.If(
            self.above_ema,
            0,
            bt.indicators.If(
                self.below_ema,
                0,
                1
            )
        )
        
        # 计算open和close与EMA的相对距离
        self.open_num = abs((self.data.open - self.ema) / self.ema)
        self.close_num = abs((self.data.close - self.ema) / self.ema)
        self.body_num = bt.indicators.If(
            self.open_num < self.close_num,
            self.open_num,
            self.close_num
        )

        # 跟踪订单和交易状态
        self.order = None
        self.trade_count = 0
        self.current_position = 0  # 0: 无仓位, 1: 多头, -1: 空头
        self.stop_price = 0  # 止损价格




    def notify_order(self, order):
        super().notify_order(order)

    def notify_trade(self, trade):
        super().notify_trade(trade)

    def next(self):
        super().next()

        if not self.position:  # 无仓位
            # 计算可用的交易资金
            available_cash = self.broker.getvalue() * self.p.risk_factor
            # 计算可买入的合约数量
            size = available_cash / self.data.close[0]
            # 做多信号
            if (self.data.close[0] > self.ema[0]) and \
               (self.upper_shadow[0] >= self.body[0] * 2) and \
               (self.cross_ema[0] == 0):
                self.order = self.buy(size=size)
                self.log(f'买入信号: 收盘价 {self.data.close[0]:.6f} > EMA10 {self.ema[0]:.6f}')
                self.stop_price = self.data.close[0] * (1 - self.p.stop_loss_pct / 100)
                # 挂单止损
                # self.order = self.sell(size=size,price=self.stop_price)

            
            # 做空信号
            elif (self.data.close[0] < self.ema[0]) and \
                 (self.lower_shadow[0] >= self.body[0] * 2) and \
                 (self.cross_ema[0] == 0):
                pass
                self.order = self.sell(size=size)
                self.log(f'卖出信号: 收盘价 {self.data.close[0]:.6f} < EMA10 {self.ema[0]:.6f}')
                self.stop_price = self.data.close[0] * (1 + self.p.stop_loss_pct / 100)
                # 挂单止损
                # self.order = self.buy(size=size,price=self.stop_price)
                
        else:  # 有仓位
            # 止损逻辑
            if self.position.size > 0:  # 多头仓位
                # stop_price = self.data.open[0] * (1 - self.p.stop_loss_pct / 100)
                # if self.data.close[0] < self.stop_price:
                if self.data.low[0] < self.stop_price:
                    pass
                    self.log(f'止损触发: 多头止损价格 {self.stop_price:.6f}')
                    self.order = self.close()  # 平多
 
                # 止盈逻辑：价格回落到前一天的EMA10均线
                elif self.data.close[0] <= self.ema[-1]:
                    self.log(f'止盈触发: 多头止盈，当日最高价格 {self.data.close[0]:.6f} <= 前一日EMA10 {self.ema[-1]:.6f}')
                    self.order = self.close()  # 平多
            
            elif self.position.size < 0:  # 空头仓位
                # stop_price = self.data.open[0] * (1 + self.p.stop_loss_pct / 100)
                # if self.data.close[0] > self.stop_price:
                if self.data.close[0] > self.stop_price:
                    self.log(f'止损触发: 空头止损价格 {self.stop_price:.6f}')
                    self.order = self.close()  # 平空
 
                # 止盈逻辑：价格回升到前一天的EMA10均线
                elif self.data.close[0] >= self.ema[-1]:
                    self.log(f'止盈触发: 空头止盈，当日最低价格 {self.data.close[0]:.6f} >= 前一日EMA10 {self.ema[-1]:.6f}')
                    self.order = self.close()  # 平空